import time

from pymodbus.client import ModbusTcpClient
from pymodbus.payload import BinaryPayloadDecoder
from pymodbus.constants import Endian


class DobotController:
    def __init__(self, ip="192.168.2.14", port=502, slave=1):
        self.client = ModbusTcpClient(host=ip, port=port)
        self.slave = slave

    # ---------- 基本连接 ----------
    def connect(self):     return self.client.connect()
    def disconnect(self):  self.client.close()

    # ---------- 触发脚本(上升沿) ----------
    def trigger_script(self, cmd: str):
        """cmd: start / pause / resume / stop / estop"""
        coils = {"start": 0, "pause": 1, "resume": 2, "stop": 3, "estop": 4}
        addr = coils.get(cmd)
        if addr is None:
            raise ValueError("无效命令，可选: start/pause/resume/stop/estop")

        # 写 True 产生上升沿
        self.client.write_coil(address=addr, value=True, slave=self.slave)
        # 马上拉回 False，为下一次触发准备
        self.client.write_coil(address=addr, value=False, slave=self.slave)

    # ---------- 运行状态 ----------
    def is_running(self) -> bool | None:
        """返回 True=运行, False=非运行, None=通讯失败"""
        r = self.client.read_discrete_inputs(address=3, count=1, slave=self.slave)
        return None if r.isError() else r.bits[0]

    # ---------- 6 轴位姿 ----------
    def read_position(self) -> dict[str, float] | None:
        """读取机械臂 6 轴位姿
        
        Returns:
            dict[str, float] | None: 包含 x,y,z,a,b,c 6个轴位置的字典，通讯失败时返回 None
        """
        start_addr, regs = 242, 12          # 6×F32 = 12 regs
        r = self.client.read_input_registers(
            address=start_addr, count=regs, slave=self.slave
        )
        if r.isError():
            return None
        
        # 使用正确的参数名称调用 convert_from_registers
        values = self.client.convert_from_registers(
            registers=r.registers,
            data_type=self.client.DATATYPE.FLOAT32,  # 使用 FLOAT32 数据类型
            word_order="little"  # 只使用 word_order 参数
        )
        axes = ["x", "y", "z", "a", "b", "c"]
        return dict(zip(axes, values))


# -------------------- Demo --------------------
if __name__ == "__main__":
    arm = DobotController()

    if not arm.connect():
        print("机械臂连接失败")
        quit()

    print("机械臂连接成功，发送 start……")
    arm.trigger_script("start")

    # 等待进入运行状态
    for _ in range(50):                    # 最多 5 秒
        run = arm.is_running()
        if run is None:
            print("读取状态失败")
            break
        if run:
            print("已进入运行状态")
            break
        time.sleep(0.1)
    else:
        print("超时仍未运行，检查 PLC 地址或模式设置")

    # 读一次位姿
    pos = arm.read_position()
    if pos:
        print("当前位置:")
        for k, v in pos.items():
            print(f"  {k.upper():>2} = {v:.3f}")

    arm.disconnect()
